package com.diandi.zuul;

import com.netflix.loadbalancer.RoundRobinRule;
import com.netflix.loadbalancer.Server;
import com.netflix.niws.loadbalancer.DiscoveryEnabledServer;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.EncodedResource;
import org.springframework.core.io.support.PropertiesLoaderUtils;

import java.io.IOException;
import java.util.*;

/**
 * 重写ribbon 分配规则 根据百分比配置
 * dd-message.192.168.3.143.8081=10
 * dd-message.192.168.3.143.8080=90
 * Created by 袁恩光 on 2018/1/8.
 */
public class DDRibbonRules extends RoundRobinRule {

    private static final String CHART_SET = "UTF-8";

    private static final String PROPERTIES_NAME = "rules.properties";

    @Override
    public Server choose(Object key) {
        Map<String, String> weights = loadRules();
        if (weights.size() == 0) {
            return super.choose(key);
        }
        try {
            List<Server> serverList = getLoadBalancer().getAllServers();
            DiscoveryEnabledServer discoveryEnabledServer = (DiscoveryEnabledServer) serverList.get(0);
            //权重..
            //服务list
            List<Integer> weightList = new ArrayList<>();
            List<String> serviceNames = new ArrayList<>();
            for (String serviceKey : weights.keySet()) {
                int indexdot = serviceKey.indexOf(".");
                String serviceId = serviceKey.substring(0, indexdot);
                String serviceName = serviceKey.substring(indexdot + 1);
                if (serviceId.equalsIgnoreCase(discoveryEnabledServer.getInstanceInfo().getAppName())) {
                    //地址和端口
                    if (Integer.valueOf(weights.get(serviceKey)) == 0) {
                        continue;
                    }
                    serviceNames.add(serviceName);
                    if (weightList.size() == 0) {
                        weightList.add(Integer.valueOf(weights.get(serviceKey)));
                    } else {
                        weightList.add(weightList.get(weightList.size() - 1) + Integer.valueOf(weights.get(serviceKey)));
                    }
                }
            }
            if (serviceNames.size() > 0) {
                //取权重随机数 并获取 index
                int index = getRandom(weightList);
                String serviceName = serviceNames.get(index);
                for (Server server : serverList) {
                    DiscoveryEnabledServer dis = (DiscoveryEnabledServer) server;
                    if ((dis.getInstanceInfo().getIPAddr() + "." + dis.getInstanceInfo().getPort()).equalsIgnoreCase(serviceName)) {
                        return server;
                    }
                }
            }
        } catch (Exception e) {
            return super.choose(key);
        }
        return super.choose(key);
    }


    /**
     * 按照权重获取对应的index
     */
    private int getRandom(List<Integer> weightList) {
        int first = weightList.get(0);
        int end = weightList.get(weightList.size() - 1);
        Random random = new Random();
        int randomindex = 0;
        while (randomindex == 0) {
            randomindex = random.nextInt(end) + first;
        }
        for (Integer integer : weightList) {
            if (randomindex <= integer) {
                randomindex = integer;
                break;
            }
        }
        return weightList.indexOf(randomindex);
    }

    private Map<String, String> loadRules() {
        try {
            Resource resource = new FileSystemResource(PROPERTIES_NAME);
            EncodedResource encodedResource = new EncodedResource(resource, CHART_SET);
            Map<Object, Object> properties = PropertiesLoaderUtils.loadProperties(encodedResource);
            Map<String, String> weights = new HashMap<>();
            for (Object obj : properties.keySet()) {
                 weights.put(obj.toString(), properties.get(obj).toString());
            }
            return weights;
        } catch (IOException e) {
            return new HashMap<>();
        }
    }
}
